# -*- coding: utf-8 -*-
import cv2 as cv
import numpy as np
from mediapipe.python.solutions import drawing_utils, selfie_segmentation

cap = cv.VideoCapture("videos/back1.mp4")

mpSelfieSegmentation = selfie_segmentation.SelfieSegmentation(model_selection=1)

BG_COLOR = (192, 192, 192) # gray
bg_image = None

while True:
    flag, frame = cap.read()

    frameRGB = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
    results = mpSelfieSegmentation.process(frameRGB)

    # Draw selfie segmentation on the background image.
    # To improve segmentation around boundaries, consider applying a joint
    # bilateral filter to "results.segmentation_mask" with "image".
    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
    print(results.segmentation_mask)

    # The background can be customized.
    #   a) Load an image (with the same width and height of the input image) to
    #      be the background, e.g., bg_image = cv2.imread('/path/to/image/file')
    #   b) Blur the input image by applying image filtering, e.g.,
    #      bg_image = cv2.GaussianBlur(image,(55,55),0)
    if bg_image is None:
        bg_image = np.zeros(frame.shape, dtype=np.uint8)
        bg_image[:] = BG_COLOR
    output_image = np.where(condition, frame, bg_image)

    cv.imshow('MediaPipe Selfie Segmentation', output_image)
    if cv.waitKey(0) & 0xFF == ord('q'):
        break

cap.release()
cv.destroyAllWindows()